# coding=utf-8

"""密码管理器
	1.有登录注册功能，每个用户有自己的密码文件来独立进行管理，登录并且添加密码后
	会吧密码保存在以自己账号命名的文件下
	2.密码分类，添加、查看密码
	3.用rsa生成公钥私钥并保存在独立的文件下，进行加密解密
	写的很乱....
"""

import pickle
import rsa


def read_message(filename):
	"""读取文件"""
	try:
		with open(filename, "rb") as f:
			passwd = pickle.load(f)
			return passwd
	except IOError:
		return False
	except Exception:
		return False


def write_message(filename, a):
	"""写入文件"""
	try:
		with open(filename, "wb") as f:
			pickle.dump(a, f)	
	except IOError:
		return False
	except Exception:
		return False
		
			
def rsa_encrypt(password):
	"""加密"""
	pupkey = read_message("rsa_pup")
	if not pupkey:
		pupkey, privkey = rsa.newkeys(512)
		write_message("rsa_pup", pupkey)
		write_message("rsa", privkey)
	temp = password.encode("utf-8")
	crypto = rsa.encrypt(temp, pupkey)
	return crypto


def rsa_decrypt(crypto):
	"""解密"""
	privkey = read_message("rsa")
	temp = rsa.decrypt(crypto, privkey)
	password = temp.decode("utf-8")
	return password


def register():
	"""注册账号"""
	print("开始注册账号...")
	username = input("输入用户名： ")
	user = read_message("user") if read_message("user") else {}
	if username in user:
		print("用户已存在！\n")
		return False
	password = input("输入用户密码： ")
	user[username] = rsa_encrypt(password)     #加密
	write_message("user", user)
	print("注册成功！\n")


def loging():
	"""登录"""
	username = input("username: ")
	temp = read_message("user")
	if not temp:
		print("用户不存在\n")
		return False
	elif username not in temp:
		print("用户不存在\n")
		return False
		
	count = 0
	while count < 3:
		passwd = input("password: ")
		if passwd != rsa_decrypt(temp.get(username)):   #解密
			print("密码错误！")
			count += 1
			continue
		print("密码正确，登录成功\n")
		break
	else:
		print("错误三次，退出\n")
	print("欢迎登录密码管理器", username)
	if not username:
		username = "empty"
	while True:
		choice = input("1:添加密码 2:查看密码 3:退出 ")
		if choice == "1":
			add_passwd(username)
		elif choice == "2":
			lookup_passwd(username)
		else:
			break


def add_passwd(file):
	"""添加密码"""
	classfy = input("输入账号密码类别：1:web 2:应用 3:银行 4:服务器 ")
	if read_message(file):
		obj = read_message(file)
	else:
		obj = {
		"web":{},
		"应用":{},
		"银行":{},
		"服务器":{}
	}
	if classfy == "1":
		username = input("账号: ")
		if username in obj.get("web"):
			print("账号密码已经存在")
			return False
		passwd = input("密码: ")
		obj.get("web")[username] = rsa_encrypt(passwd)
	elif classfy == "2":
		username = input("账号: ")
		if username in obj.get("应用"):
			print("账号密码已经存在")
			return False
		passwd = input("密码: ")
		obj.get("应用")[username] = rsa_encrypt(passwd)
	elif classfy == "3":
		username = input("账号: ")
		if username in obj.get("银行"):
			print("账号密码已经存在")
			return False
		passwd = input("密码: ")
		obj.get("银行")[username] = rsa_encrypt(passwd)
	elif classfy == "4":
		username = input("账号: ")
		if username in obj.get("服务器"):
			print("账号密码已经存在")
			return False
		passwd = input("密码: ")
		obj.get("服务器")[username] = rsa_encrypt(passwd)
	write_message(file, obj)


def lookup_passwd(file):
	"""查看密码"""
	obj = read_message(file)
	if not obj:
		print("没有密码文件\n")
		return False
	a = {"1":"web", "2":"应用", "3":"银行", "4":"服务器"}
	for key, item in a.items():
		print(key, item)
	classfy = input("输入你的选择: ")
	if classfy == "1":
		if not obj.get("web"):
			print("There is nothing")
			return False
		for key in obj.get("web").keys():
			print("已经保存的账号：", key)
		u = input("输入查看账号: ")
		if u not in obj.get("web"):
			print("不存在账号")
			return False
		print(rsa_decrypt(obj.get("web")[u]))
	elif classfy == "2":
		if not obj.get("应用"):
			print("There is nothing")
			return False
		for key in obj.get("应用").keys():
			print("已经保存的账号：" + key)
		u = input("输入查看账号: ")
		if u not in obj.get("应用"):
			print("不存在账号")
			return False
		print(rsa_decrypt(obj.get("应用")[u]))
	elif classfy == "3":
		if not obj.get("银行"):
			print("There is nothing")
			return False
		for key in obj.get("银行").keys():
			print("已经保存的账号：" + key)
		u = input("输入查看账号: ")
		if u not in obj.get("银行"):
			print("不存在账号")
			return False
		print(rsa_decrypt(obj.get("银行")[u]))
	elif classfy == "4":
		if not obj.get("服务器"):
			print("There is nothing")
			return False
		for key in obj.get("服务器").keys():
			print("已经保存的账号：" + key)
		u = input("输入查看账号: ")
		if u not in obj.get("服务器"):
			print("不存在账号")
			return False
		print(rsa_decrypt(obj.get("服务器")[u]))
	else:
		return False


def main():
	while True:
		print("密码管理器：\n")
		choise = input("1:登录 2:注册 3:退出 : ")
		if choise == "1":
			loging()
		elif choise == "2":
			register()
		else:
			print("\n退出")
			break


if __name__ == "__main__":
	main()
